package com.quwan.utils;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

@Slf4j
public class ThreadPoolUtils {
    private static class Pool{
        private static final ThreadPoolUtils instance = new ThreadPoolUtils();
    }

    private static ThreadPoolExecutor executor;

    private ThreadPoolUtils() {
    }

    public static ThreadPoolUtils getInstance() {
        return Pool.instance;
    }

    public ThreadPoolExecutor getPool() {
        if (executor == null) {
            executor = MdcThreadPoolExecutor.newWithInheritedMdc(
                    "IoThread",
                    10,
                    20,
                    2,
                    TimeUnit.SECONDS,
                    new LinkedBlockingQueue<>(50)
            );
        }
        return executor;
    }
}

@Slf4j
class MdcThreadPoolExecutor extends ThreadPoolExecutor {
    @Getter
    final private String name;
    final private boolean useFixedContext;
    final private Map<String, String> fixedContext;

    /**
     * 从提交线程获取MDC。
     */
    public static MdcThreadPoolExecutor newWithInheritedMdc(String name, int corePoolSize, int maximumPoolSize, long keepAliveTime,
                                                            TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        return new MdcThreadPoolExecutor(name, null, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    /**
     * 从创建池的线程中获取固定MDC，后续该线程池中所有线程都是该MDC
     */
    public static MdcThreadPoolExecutor newWithCurrentMdc(String name, int corePoolSize, int maximumPoolSize, long keepAliveTime,
                                                          TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        return new MdcThreadPoolExecutor(name, MDC.getCopyOfContextMap(), corePoolSize, maximumPoolSize, keepAliveTime, unit,
                workQueue);
    }

    /**
     * 指定线程池所有线程的固定MDC
     */
    public static MdcThreadPoolExecutor newWithFixedMdc(String name, Map<String, String> fixedContext, int corePoolSize,
                                                        int maximumPoolSize, long keepAliveTime, TimeUnit unit,
                                                        BlockingQueue<Runnable> workQueue) {
        return new MdcThreadPoolExecutor(name, fixedContext, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    private MdcThreadPoolExecutor(String name, Map<String, String> fixedContext, int corePoolSize, int maximumPoolSize,
                                  long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
        this.fixedContext = fixedContext;
        useFixedContext = (fixedContext != null);
        this.name = name;
    }

    @Override
    protected void terminated() {
        super.terminated();
        log.info("MDC线程池结束");
    }

    private Map<String, String> getContextForTask() {
        return useFixedContext ? fixedContext : MDC.getCopyOfContextMap();
    }

    @Override
    public void execute(Runnable command) {
        log.info("线程池[{}]: 线程数(50):{}, 峰值线程数：{}, 活动线程数:{}, 队列数(200):{}, 执行任务数:{}, 完成任务数:{}",
                "IoThread",
                getPoolSize(),
                getLargestPoolSize(),
                getActiveCount(),
                getQueue().size(),
                getTaskCount(),
                getCompletedTaskCount()
        );
        super.execute(wrap(command, getContextForTask(), RequestContextHolder.currentRequestAttributes()));
    }

    private static Runnable wrap(final Runnable runnable, final Map<String, String> context, RequestAttributes requestAttributes) {
        return () -> {
            Map<String, String> previous = MDC.getCopyOfContextMap();
            RequestContextHolder.setRequestAttributes(requestAttributes, true);
            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            try {
                runnable.run();
            } finally {
                if (previous == null) {
                    MDC.clear();
                } else {
                    MDC.setContextMap(previous);
                }
            }
        };
    }
}

